package middleware

import (
	"bytes"
	"encoding/json"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
	"io"
	"net/http"
	"strings"
	"time"
)

// ResponseRecorder 用于捕获响应
type ResponseRecorder struct {
	http.ResponseWriter
	Body       *bytes.Buffer
	StatusCode int
}

func NewResponseRecorder(w http.ResponseWriter) *ResponseRecorder {
	return &ResponseRecorder{
		ResponseWriter: w,
		Body:           &bytes.Buffer{},
		StatusCode:     http.StatusOK,
	}
}

func (r *ResponseRecorder) Write(b []byte) (int, error) {
	r.Body.Write(b)
	return r.ResponseWriter.Write(b)
}

func (r *ResponseRecorder) WriteHeader(statusCode int) {
	r.StatusCode = statusCode
	r.ResponseWriter.WriteHeader(statusCode)
}

func LoggingRequestMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// 记录请求开始时间
		startTime := time.Now().Local()

		body := []byte("")
		headerContentType := r.Header.Get("Content-Type")
		if headerContentType != "" && !strings.HasPrefix(headerContentType, "multipart/form-data") {
			bodyJson, err := io.ReadAll(r.Body)
			// 读取请求体
			if err != nil {
				logx.Errorf("读取请求体失败: %v", err)
				httpx.Error(w, err)
				return
			}
			// 恢复请求体，以便后续处理
			r.Body = io.NopCloser(bytes.NewBuffer(bodyJson))
			body = bodyJson
		}

		// 记录请求日志
		replacer := strings.NewReplacer("\n", "", "\t", "", "    ", "")
		cleanBody := replacer.Replace(string(body))
		logx.WithContext(r.Context()).Infof("请求方法: %s, 请求路径: %s, 请求体: %s", r.Method, r.URL.Path, cleanBody)
		// 使用自定义的 ResponseRecorder
		recorder := NewResponseRecorder(w)
		next(recorder, r)

		staffIdRaw := r.Context().Value("sid")
		staffId := int64(0)
		if _, ok := staffIdRaw.(json.Number); ok {
			staffId, _ = staffIdRaw.(json.Number).Int64()
		}
		// 调用下一个处理器
		//next.ServeHTTP(w, r)
		responseBody := recorder.Body.String()
		logx.WithContext(r.Context()).Infof("登录staffId：%d,响应状态码: %d, 响应体: %s, 处理时间: %v", staffId, recorder.StatusCode, responseBody, time.Since(startTime))
	})
}
